package sharding_DB

import (
	"context"
	"database/sql"
	db2 "gitee.com/go-web-2/orm_framework/session/db"
	sql2 "gitee.com/go-web-2/orm_framework/sql"
	"math/rand"
)

type MasterSlavesDB struct {
	master *db2.DB
	slaves []*db2.DB
	db2.Core
}

func (db *MasterSlavesDB) IsDB() {

}

func (db *MasterSlavesDB) ExecContext(ctx context.Context, query string, args ...any) *sql2.QueryResult {
	res, err := db.master.ExecContext(ctx, query, args...)
	return &sql2.QueryResult{
		Res: res,
		Err: err,
	}
}

func (db *MasterSlavesDB) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
	if db.slaves == nil || len(db.slaves) == 0 {
		return db.master.QueryContext(ctx, query, args...)
	}

	idx := rand.Intn(len(db.slaves)) // 负载均衡：随机数；可以完善从库注册与发现和负载均衡方法
	return db.slaves[idx].QueryContext(ctx, query, args...)
}

func (db *MasterSlavesDB) BeginTx(ctx context.Context, txOpts *sql.TxOptions) (*db2.Tx, error) {
	tx, err := db.master.Db.BeginTx(ctx, txOpts)
	if err != nil {
		return nil, err
	}
	return &db2.Tx{
		Tx: tx,
		Db: &db2.DB{
			Db:   db.master.Db,
			Core: db.Core,
		},
	}, nil
}

func NewMasterSlavesDB(master *db2.DB, slaves []*db2.DB) *MasterSlavesDB {
	db := &MasterSlavesDB{
		master: master,
		slaves: slaves,
		Core:   master.Core,
	}

	return db
}
